function [BETA,R,varargout]=regression2(Y,X,varargin)
%{
This function executes either OLS or IV regressions
It allows for multiple categorical variables to be differenced out or used as dummies.

Inputs (Y,X,E,Z,FE,ops):
  Required:
  Y(n-1) - Dependent Variable
  X(n-k) - Exogenous Independent Regressors
  Optional
  E(n-p) - Endogenous Variables
  Z(n-q) - Instruments for endogenous variables; must have at least as many instruments as endogenous variables
  FE(n-m) - Categorical variables to demean or include as dummies; program will demean on the first variable 
  ops(struct) - see below

Outputs: BETA, SE, VCV, Y_HAT, XB_HAT, G_HAT

Notes:
(1) Scaling and Recentering: The program scales all variables to have a std. dev of 1 and a mean value of zero. This gives stability to the
matrices and allows for better estimation of which vectors may be collinear. However, it affects the standard deviation of the constant, as
it is estimated differently than under the usual regression techniques. 
%}
%#ok<*MINV>

%% A) SETUP
% Define default options
def = struct(...
  'vce','standard',... Standard errors; options are 'standard','robust','cluster'
  'cluster',[],...     Variable on which to cluster; 'cluster' must be specificed for vce
  'weights',[],...     Variable to use as weights in regression
  'scale',1,...        0/1 indicator for the use of scaling and recentering
  'exovars',{{}},...   List of names for exogenous variables (used in printing results)
  'endvars',{{}},...   List of names for endogenous variables (used in printing results)
  'display','small'); %Display estimates of all variables('large'), only X and Xe variables ('small') or 'none'

% Unpack varargin
if nargin>=3, E =varargin{1}; else E=[]; end;
if nargin>=4, Z =varargin{2}; else Z=[]; end;
if nargin>=5, FE=varargin{3}; else FE=[]; end;
if nargin>=6, ops=varargin{4}; else ops=def; end;

% Get options
t=fieldnames(def); for i=1:numel(t), if ~isfield(ops,t{i}), ops.(t{i})=def.(t{i}); end; end; clear i; 

% Check for errors
  % Y or X not vectors
  if ~isvector(Y) || size(Y,2)~=1, error('MATLAB:regression:Y','First argument (Y) must be a col vector'); end;
  if ~ismatrix(X), error('MATLAB:regression:X','Second argument (X) must be a matrix'); end;
  n=length(Y); if size(X,1)~=n, error('MATLAB:regression:X','First and Second arguments have different no. of observations'); end;
  % Options correctly specified
  if ~iscellstr(ops.('exovars')) || ~iscellstr(ops.('endvars')), error('MATLAB:regression:ops','Variable list names are not cell arrays of strings'); end;
  t=ops.('weights'); if ~isempty(t) && ~isvector(t), error('MATLAB:regression:ops','Option ''weights'' is not a vector'); end;
  if ~isempty(t) && any(size(t)~=[n 1]), error('MATLAB:regression:ops','Option ''weights'' does not have same number of observations as Y'); end;
  t=ops.('vce'); if ~ischar(t), error('MATLAB:regression:ops','Option ''vce'' incorrectly specified'); end;
  m=ops.('cluster'); if strcmpi(t,'cluster') && (~isvector(m) || any(size(m)~=[n 1])), error('MATLAB:regression:ops','Option ''cluster'' incorrectly specified'); end;
  if ~any(strcmpi(t,{'cluster','standard','robust'})), error('MATLAB:regression:ops','Option ''vce'' incorrectly specified'); end;
  if ~ischar(ops.('display')), error('MATLAB:regression:ops','Option ''display'' incorrectly specified'); end;
  t=ops.('scale'); if ~isfinite(t) || ~isscalar(t), error('MATLAB:regression:scale','Option ''scale'' incorrectly specified'); end;
  % Endogenous Vars and Instruments
  if ~ismatrix(E), error('MATLAB:regression:E','Third argument incorrectly specified'); end;
  if ~ismatrix(Z), error('MATLAB:regression:Z','Fourth argument incorrectly specified'); end;
  if ~isempty(E) && size(E,1)~=n, error('MATLAB:regression:E','Third argument (E) has different number of observations than Y'); end;
  if ~isempty(Z) && size(Z,1)~=n, error('MATLAB:regression:Z','Fourth argument (Z) has different number of observations than Y'); end;
  if size(Z,2)<size(E,2), error('MATLAB:regression:Z','Not enough instruments specified'); end;
  % Fixed Effects
  if ~ismatrix(FE), error('MATLAB:regression:FE','Fifth argument incorrectly specified'); end;
  if ~isempty(FE) && size(FE,1)~=n, error('MATLAB:regression:FE','Fifth argument (FE) has different number of observations than Y'); end;

% Input option variables into variables and exand empty variables to appropriate size
W=ops.('weights'); if strcmpi(ops.('vce'),'cluster'), C=ops.('cluster'); else C=NaN(n,0); end;
if isempty(W), W=ones(n,1); weights=false; else weights=true; end; 
if isempty(C), C=NaN(n,0); end; 
if isempty(Z), Z=NaN(n,0); end; 
if isempty(E), E=NaN(n,0); end; 
if isempty(FE), FE=NaN(n,0); end;
if any(strcmpi(ops.('vce'),{'cluster','robust'})), vce=true; else vce=false; end;

%% B) PREPARE DATA FOR REGRESSIONS

% Remove missing data
t=[Y X E Z FE W C]; OBS=all(isfinite(t),2); v=~OBS; clear t; 
Y(v,:)=[]; X(v,:)=[]; E(v,:)=[]; Z(v,:)=[]; FE(v,:)=[]; W(v,:)=[]; C(v,:)=[]; 

% Get number of variables and number of observations
N=length(Y); Kx=size(X,2); Ke=size(E,2); Kz=size(Z,2);

% Make weights add up to N
if weights, W=W*N/sum(W); end;

% Create lists of varnames for X and E
xl = ops.('exovars'); t=length(xl); for i=(t+1):Kx, xl{i}=['x' int2str(i)]; end;
el = ops.('endvars'); t=length(el); for i=(t+1):Ke, el{i}=['e' int2str(i)]; end; for i=1:length(el), el{i}=['*' el{i}]; end;

% Scale and Recenter for stability (weighted means, weighted std dev)
if ops.('scale')
  MuZ=(W'*Z)/N; MuE=(W'*E)/N; MuX=(W'*X)/N;
  SclX=sqrt(W'*(bsxfun(@plus,X,-MuX).^2)/N); SclX(SclX==0)=1; 
  SclE=sqrt(W'*(bsxfun(@plus,E,-MuE).^2)/N); SclE(SclE==0)=1; 
  SclZ=sqrt(W'*(bsxfun(@plus,Z,-MuZ).^2)/N); SclZ(SclZ==0)=1; 
  YEXZ=[Y bsxfun(@rdivide,bsxfun(@plus,E,-MuE),SclE) ...
          bsxfun(@rdivide,bsxfun(@plus,X,-MuX),SclX) ...
          bsxfun(@rdivide,bsxfun(@plus,Z,-MuZ),SclZ)]; 
else YEXZ = [Y E X Z]; 
end;

% Create vectors that indicate location of variables in YEXZ
yv=1; j=1; ev=(1:Ke)+1; xv=(1:Kx)+max([yv ev]); zv=(1:Kz)+max([yv ev xv]); 

% Create Dummies for Categorical Variables
  % Get unique values of each FE
  f=size(FE,2); uFE=cell(f,1); t=NaN(f,1);
  for i=1:f, uFE{i}=unique(FE(:,i)); t(i)=length(uFE{i}); end; 
  % Find FE to demean on (the variables with most distinct groups)
  if f, iFw=find(max(t)==t,1); else iFw=0; end;
  % Preallocate, including constant
  Kf=1; for i=1:length(uFE), if i~=iFw, Kf=Kf+length(uFE{i})-1; end; end;
  xF=ones(N,Kf); fl=cell(1,Kf); fl{1}='cons'; v=2;  
  % Create dummies and list of names for dummies
  for i=1:f, if i==iFw, continue; end; t=['F' int2str(i) '_'];
    for j=2:length(uFE{i}), xF(:,v)=FE(:,i)==uFE{i}(j); fl{v}=[t int2str(j)]; v=v+1; end;
  end; clear v; 

% Concatenate FE 
YEXZ=[YEXZ xF]; fv=(1:Kf)+max([yv ev xv zv]);

% Demean variables of FE group (and store if predicting outcomes)
Kw=0; 
if nargout>3 && iFw, Y_bar=NaN(Kw,length([yv ev xv fv])); end;
if iFw, u=uFE{iFw}; Kw=length(u); tt=(W'*YEXZ)./N;
  for i=1:Kw, j=FE(:,iFw)==u(i); 
    if ~weights, t=mean(YEXZ(j,:),1); else t=(W(j)'*YEXZ(j,:))./sum(W(j)); end;
    YEXZ(j,:)=bsxfun(@plus,YEXZ(j,:),tt-t); 
    % Store Y_bar and X_bar for predicting within estimates below
    if nargout>3, p=t-tt; Y_bar(i,:)=p([yv ev xv fv]); end;
  end;
end;
% Scale Y_bar back to original size if scalling is applied
if ops.('scale') && nargout>3 && iFw, s=[SclE SclX]; t=1+(1:length(s)); Y_bar(:,t)=bsxfun(@times,Y_bar(:,t),s); end;


%% C) REGRESSION (see Greene [5th ed], Theorem 10.6)

% Initiate Beta, SE, Var Names, and Identifier
BETA=zeros(Ke+Kx+Kf,1); Bid=1:length(BETA);

% Remove collinear cols (std=0 [except for constant], this also removes all zeros)
s=std(YEXZ,[],1)<2*N*eps; s(fv(1))=false; 
t=s([ev xv fv]); Bid(t)=[]; 
ev(s(ev))=[]; xv(s(xv))=[]; zv(s(zv))=[]; fv(s(fv))=[]; % Reduce variables
Ke=length(ev); Kx=length(xv); Kz=length(zv); Kf=length(fv); 

% Square matrix
if ~weights, XX=YEXZ'*YEXZ; else XX=bsxfun(@times,YEXZ,W)'*YEXZ; end;

% Solve for 'b', 'x', and 'ixz' under OLS, IV EXACT, and IV OVERIDENTIFIED
x=[ev xv fv]; z=[zv xv fv]; 
if ~Ke || Ke==Kz, % OLS or IV Exact : (Z'X)^-1 (Z'Y)
  xb=base(XX(z,x),1:(Ke+Kx)); x=x(xb); z=z(xb);
  ixz=inv(XX(z,x)); b=ixz*XX(z,yv);  
else % IV Overidentified: (Xh'Xh)^-1 (Xh'Y) ; Xh = Z(Z'Z)^-1 (Z'*X)
  zb=base(XX(z,z)); z=z(zb); 
  xzizz=XX(x,z)/XX(z,z); xhxh=xzizz*XX(z,x); 
  xb=base(xhxh,1:(Ke+Kx)); xhxh=xhxh(xb,xb); x=x(xb); xzizz=xzizz(xb,:);
  ixz=inv(xhxh); b=ixz*xzizz*XX(z,yv); 
end;

% Solve for residuals 'r' 
R=YEXZ(:,yv) - YEXZ(:,x)*b; r2w=(R.^2).*W; 

% Get basic statistics
r=W'*Y/N-Y; SSRy=(r.*W)'*r;                          % SSR for Y on 1, no demeaning
SSR=sum(r2w); SSRw = XX(yv,yv) - XX(yv,fv(1))^2 / N; % SSR for y on 1, w/demeaning: ( Y'Y - (Y'1)^2 /N )
K=length(x); 
DOF=N-K-Kw+(Kw>0);  % (Kw>0) in DOF is because of 1 group of within FE not being used
ADJ = N/DOF;
R2 = 1-SSR./SSRy;  aR2 = 1 - (N-1)./(N-K-Kw).*(1-R2);
R2w = 1-SSR./SSRw; aR2w= 1 - (N-1)./(N-K).*(1-R2w);

% Get vcv matrix
s2 = SSR/N*ADJ;
if ~vce 
  if ~Ke || Ke<Kz, vcv=ixz * s2;    % OLS or IV OVERID
  else vcv=(ixz*XX(z,z)*ixz') * s2; % IV EXACT
  end;
else % Robust or Cluster: 
  % (A) Get weigting matrix Z*W*Z
  if strcmpi(ops.('vce'),'robust'), 
    zwz=bsxfun(@times,YEXZ(:,z),r2w)'*YEXZ(:,z);
  else % Cluster
    zwz=zeros(length(z)); t=unique(C); ng=length(t); if weights, rw=R.*sqrt(W); else rw=R; end;
    for i=1:ng, j=t(i)==C; p=YEXZ(j,z)'*rw(j); zwz=zwz+p*(p'); end;
    ADJ=ADJ*ng/(ng-1)*(N-1)/N;
  end;
  % (B) Calculate VCV
  if Ke<Kz, p=ixz*xzizz; vcv=p*zwz*p' * ADJ; % IV OVERID
  else vcv=ixz*zwz*ixz * ADJ;                % OLS or IV EXACT
  end;
end;

%% D) RETURN RESULTS

% Store 'b' and 'vcv' in appropriate places
t=Bid(xb); BETA(t)=b; VCV=zeros(length(BETA)); VCV(t,t)=vcv; 

% Re-scale and recenter (variance of constant given by Var(a0)=Var(a0_hat - x_bar*b_hat) )
if ops.('scale')
  s=[SclE SclX]'; t=length(s); BETA(1:t)=BETA(1:t)./s; VCV(1:t,1:t)=VCV(1:t,1:t)./(s*s'); 
  p=[MuE MuX]; BETA(t+1)=BETA(t+1)-p*BETA(1:t);  
  j=t+1; p=[-p 1]; v0=p*VCV(1:j,1:j)*p'; % V(sum{a*X})=sum{sum{ai*aj*Cov(Xi,Xj)}}
  v1=VCV(1:t,1:j)*p';                    % Cov[sum{a*X},Z] = sum{a*Cov[X,Z]}
  VCV(j,j)=v0; VCV(1:t,j)=v1; VCV(j,1:t)=v1';
end;

% Input NaN in VCV when regressor was dropped
t=NaN(size(BETA));t(Bid(xb))=0;t=isnan(t); VCV(t,:)=NaN; VCV(:,t)=NaN; 
% Return SE and VCV if solicited
SE=sqrt(diag(VCV)); if nargout>=2 , varargout{1}=SE; end; if nargout>=3, varargout{2}=VCV; end;

% Predict Values
if nargout>3
  % Initialize variables
  n=length(OBS); xb_hat=NaN(n,1); g_hat=NaN(n,length(uFE));
  % y_hat (need to add wFE estimates later)
  t=size(E,2)+size(X,2); xb_hat(OBS)= [E X]*BETA(1:t)+BETA(t+1); 
  % wFE estimates ('d' in Stata; d=y_bar - x_bar*b_hat)
  if iFw, d=NaN(N,1); u=uFE{iFw}; p=Y_bar*[1;-BETA]; for i=1:Kw, d(FE(:,iFw)==u(i))=p(i); end; end; 
  % G_hat (dummy FEs)
  tx=1; tb=j(end); 
  for i=1:length(uFE), % Cycle over FE groups
    if i==iFw, g_hat(OBS,i)=d; continue; end; 
    v=1:(length(uFE{i})-1);  % FEs within a FE group, omitting 1st category
    g_hat(OBS,i)=xF(:,tx+v)*BETA(tb+v); tx=tx+v(end); tb=tb+v(end); 
  end;
  % Xb_hat, includes the constant
  y_hat=xb_hat+sum(g_hat,2);
  varargout{3}=y_hat; if nargout>4, varargout{4}=xb_hat; end; if nargout>5, varargout{5}=g_hat; end;
end;

%% E) PRINT RESULTS
if ~strcmpi(ops.('display'),'none') 
  % Header
  fprintf(1,'\n___________________________________________________\n ');
  fprintf(1,'\n   '); if Ke, fprintf(1,'I V'); else fprintf(1,'O L S'); end; fprintf(1,'   R E G R E S S I O N   R E S U L T S');
  fprintf(1,'\n___________________________________________________');
  % Estimates
  fprintf(1,'\n%-10s %10s %10s %10s \n','VARIABLE','ESTIMATE','STD. ERR','T-STAT');
  if strcmpi(ops.('display'),'small'), vl=[el xl fl{1}]; else vl=[el xl fl]; end;
  for i=1:length(vl), fprintf(1,'%-10s %10.3f %10.3f %10.3f \n',vl{i},BETA(i),SE(i),abs(BETA(i)/SE(i))); end;
  if Ke, fprintf(1,'________\n* mark endogenous variables\n'); end;
  % Constants
  fprintf(1,'\nREGRESSION CONSTANTS:\n');
  fprintf(1,'%2s %-7u \t %4s %-7u','N=',N,'DOF=',DOF);
  if Kw, fprintf('\nWithin Groups: %-7u',Kw); end; if Kf>1, fprintf('\t FE dummies: %-7u',Kf-1); end;
  fprintf(1,'\n%19s %-5.3f \t %-5s %-5.3f','Overall R2s: unadj.=',R2,'adj.=',aR2);
  if Kw, fprintf(1,'\n%19s %-5.3f \t %-5s %-5.3f','Within R2s:  unadj.=',R2w,'adj.=',aR2w);end;
  % Notes
  fprintf(1,'\n\nNOTES:');
  if strcmpi(ops.('vce'),'robust'), fprintf('\n* Heteroskedastic robust standard errors'); end;
  if strcmpi(ops.('vce'),'cluster'), fprintf('\n* Clustered standard errors'); end;
  if Kw, fprintf('\n* Within estimator used to demean fixed effects'); end;
  if Kf>1 && strcmpi(ops.('display'),'small'), fprintf('\n* Additional fixed effects included but not shown'); end;
  if weights, fprintf(1,'\n* Weighted regression'); end;
  if ops.('scale'), fprintf(1,'\n* Std. Err. on constant may differ from other \n  algorithms because of recentering. However, it\n  is consistent'); end;
  fprintf(1,'\n___________________________________________________ \n');
end;

end